import os
import sys
import subprocess
import json
from tqdm import tqdm
from pathlib import Path
def color_print(msg=""):
        print_msg = "\033[32m"
        print_msg += msg
        print_msg += "\033[0m"
        print(print_msg)

# 默认配置选项
DEFAULT_CONFIG = {
    "IRGEN_PATH": "/Talon_tools/IRGen/irgen_entry.py",
    "MAPLE_ROOT":"/Talon_tools/IRGen/Maple/OpenArkCompiler",
    "MAPLE_CLANG":"/Talon_tools/IRGen/Maple/OpenArkCompiler/tools/bin/clang",
    "MAPLE_CLANGCPP":"/Talon_tools/IRGen/Maple/OpenArkCompiler/tools/bin/clang++",
    "MAPLE_BIN":"/Talon_tools/IRGen/Maple/OpenArkCompiler/output/aarch64-clang-release/bin/hir2mpl",
    "MAPLE_GCC":"/Talon_tools/IRGen/Maple/OpenArkCompiler/output/tools/gcc-linaro-7.5.0/bin/aarch64-linux-gnu-gcc",
    "MAPLE_MAPLE":"/Talon_tools/IRGen/Maple/OpenArkCompiler/output/aarch64-clang-release/bin/maple",
    "talon_executable": "/TalonInArkCompiler/output/aarch64-clang-release/bin/maple",
    "talon_bool_options": {
        "talon": True,
        "debug-talon": True,
        "no-inline": True,
        "psa-enable-side-effect-source": True,
        "psa-enable-arg-symbol":True,
        "enable-heap-alloc-failure": False,
        "enable-fd-alloc-failure": False
    },
    "checker_options":{
        "ps-npd": True,
        "ps-dbz": True,
        "ps-dbf": True,
        "ps-rsa": False,
        "ps-fnhm": False,
        "ps-uuv": False,
        "ps-ml": False,
        "ps-fduac": False,
        "ps-bcns": False,
        "ps-fdl": False,
        "ps-fdl2": False,
        "ps-bbsf": False
    },
    "talon_value_options":{
        "restrict-cg-size": 1,
        "inline-depth": 6,
        "nworkers": 20,
        "report-format": "json",
        "report-pass-line": 1,
        "execution-mode": "normal",
        "system-timeout":7200
    },
    "report_suffix": "bugreport.json"
    
}

CONFIG_FILE = 'config.json'

def get_absolute_path(user_input):
    path=Path(user_input).resolve
    return path
def run_command(command=[], env={}, cwd=""):
    if not env:# 环境变量
        env = dict(os.environ)
    if not cwd:# 调用时的环境
        cwd = os.getcwd()
    print(f"{cwd}")
    ret = subprocess.call(args=command, shell=False, env=env, cwd=cwd)
    return ret

# 保存默认配置文件
def save_default_config():
    with open(CONFIG_FILE, 'w') as f:
        json.dump(DEFAULT_CONFIG, f, indent=4)
    print(f"默认配置文件已保存为 {CONFIG_FILE}")

# 读取配置文件
def load_config():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path=os.path.join(script_dir,CONFIG_FILE) 
    if not os.path.exists(config_path):
        save_default_config()
    
    with open(config_path, 'r') as f:
        return json.load(f)

# 生成 IR
def generate_ir(config,project_path,build_command):
    irgen_entry_bash=os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),"IRGen/irgen_entry.py")
    os.chdir(project_path)
    # print(f"build_command {build_command}")
    prefix=["python3", irgen_entry_bash,
            "--specify-maple-root",config['MAPLE_ROOT'],
            "--specify-maple-maple",config['MAPLE_MAPLE'],
            "--specify-maple-gcc",config['MAPLE_GCC'],
            "--specify-maple-bin",config['MAPLE_BIN'],
            "--specify-maple-clangcpp",config['MAPLE_CLANGCPP'],
            "--specify-maple-clang",config['MAPLE_CLANG'],
              "--"]
    irgen_command=prefix + build_command
    # print(f"irgen_command {irgen_command}")
    subprocess.run(irgen_command, check=True)
    print("Generation Of MAPLE IR Is Complete")

# 将配置文件转化成运行参数
def generate_options(config):
    options="".join([
        f"--{key} " for key, value in config['talon_bool_options'].items() if value
    ])
    options+="".join([
        f"--{key} " for key, value in config['checker_options'].items() if value
    ])
    options+=" ".join([
        f"--{key}={value}" for key, value in config['talon_value_options'].items() 
    ])
    print(f"options:{options} \n")
    return options


# 检测 IR 中的漏洞
def detect_vulnerabilities(project_path, config):
    color_print("Talon: a tool for detecting various Vulnerablities")
    ir_path = os.path.join(project_path, '.IRGEN', 'top')
    if not os.path.exists(ir_path):
        color_print("[ERROR] Folder Is Not Exist")
        return
    options = generate_options(config)
    report_suffix = config['report_suffix']
    talon_executable=config['talon_executable']
    bug_report_output_folder=os.path.join(project_path,'bug_reports')
    os.makedirs(bug_report_output_folder, exist_ok=True)
    log_folder=os.path.join(bug_report_output_folder,'log')
    os.makedirs(log_folder, exist_ok=True)
    error_log_folder=os.path.join(bug_report_output_folder,'error')
    os.makedirs(error_log_folder, exist_ok=True)
    # 获取所有 .mpl 文件
    mpl_files = [filename for filename in os.listdir(ir_path) 
                 if filename.endswith('.mpl')]
    with tqdm(mpl_files, unit='file') as pbar:
        for filename in pbar:
            pbar.set_description("Processing %s" % filename) 
            name_without_extension = os.path.splitext(filename)[0]
            mpl_file = os.path.join(ir_path, filename)

            target_report_filename=name_without_extension+'_'+(report_suffix)
            report_file_path= os.path.join(bug_report_output_folder,target_report_filename)

            command = f"{talon_executable} --run=mpl2mpl --option='-O2 {options} --report={report_file_path}' --infile {mpl_file}"
           
            output_log=os.path.join(log_folder,name_without_extension+'.log')
            error_log=os.path.join(error_log_folder,name_without_extension+'_error.log')
            run_command(command,output_log,error_log)
            # result=subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True,text=True )
            
            # with open(output_log, 'w') as output_log:
            #     # print(f"talon command:{command}")
            #     output_log.write(f"talon command:{command}")
            #     output_log.flush()
            #     output_log.write(result.stdout)
            #     output_log.flush()
            # if result.stderr:
            #     with open(error_log, 'w') as error_log:
            #         error_log.write(result.stderr)
            #         error_log.flush()
    total_bug_path=os.path.join(bug_report_output_folder,"total_bugs.json")
    color_print(f"[INFO] Vulnerability Detection Is Complete")
    color_print(f"[INFO] Results Are Saved In {total_bug_path}")
    

    script_dir = os.path.dirname(os.path.abspath(__file__))
    remove_dumplicate_bash=os.path.join(script_dir,"handle_bugreport/remove_dumplicate_bugreport.py")
    subprocess.call(['python3',remove_dumplicate_bash,bug_report_output_folder])

# 运行命令并捕获异常
def run_command(command, output_log_path, error_log_path):
    try:
        with open(output_log_path, 'w') as output_log:
            # 写入命令到输出日志
            output_log.write(f"talon command:{command}\n")
            output_log.flush()

            # 运行命令并捕获 stdout 和 stderr
            result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True, text=True)

            # 将 stdout 写入输出日志
            output_log.write(result.stdout)
            output_log.flush()

            # 检查是否有 stderr 输出
            if result.stderr:
                with open(error_log_path, 'w') as error_log:
                    error_log.write(result.stderr)
                    error_log.flush()

    except subprocess.CalledProcessError as e:
        with open(output_log_path, 'w') as output_log:
            output_log.write(f"talon command:{command}\n")
            output_log.write(f"Command failed with return code {e.returncode}\n")
            output_log.write(e.stdout)
            output_log.flush()

        with open(error_log_path, 'w') as error_log:
            error_log.write(e.stderr)
            error_log.flush()

    except Exception as e:
        with open(output_log_path, 'w') as output_log:
            output_log.write(f"talon command:{command}\n")
            output_log.write(f"An unexpected error occurred: {e}\n")
            output_log.flush()

        with open(error_log_path, 'w') as error_log:
            error_log.write(f"An unexpected error occurred: {e}\n")
            error_log.flush()
def get_build_command():
    args = sys.argv
    index = args.index("--")
    build_args = args[index+1:]
    return build_args
def main():
    build_command = get_build_command()
    project_path = os.getcwd()
    config = load_config()

    project_path=os.path.abspath(project_path)

    generate_ir(config, project_path, build_command)
    detect_vulnerabilities(project_path, config)
def test():
    script_dir =(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 
    print(f"当前脚本所在的root路径: {script_dir}")
    
    # ret = subprocess.call(args=command, shell=True, env=env, cwd=cwd)

if __name__ == "__main__":
    # test()
    # exit()
    main()